{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "86c20c93",
   "metadata": {},
   "source": [
    "# 测试 profiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "534ab1fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "import tvm.testing\n",
    "from tvm.runtime.vm import VirtualMachine\n",
    "from tvm import relax, rpc\n",
    "from tvm.contrib import utils\n",
    "from tvm.relax.testing import nn\n",
    "from tvm.script import relax as R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c76e756e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_exec(data_shape):\n",
    "    builder = relax.BlockBuilder()\n",
    "    weight1_np = np.random.randn(64, 64).astype(\"float32\")\n",
    "    weight2_np = np.random.randn(64, 64).astype(\"float32\")\n",
    "\n",
    "    with builder.function(\"main\"):\n",
    "        model = nn.Sequential(\n",
    "            nn.Linear(data_shape[1], weight1_np.shape[0], bias=False),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(weight2_np.shape[0], weight2_np.shape[1], bias=False),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "        data = nn.Placeholder(data_shape, name=\"data\")\n",
    "        output = model(data)\n",
    "        params = [data] + model.parameters()\n",
    "        builder.emit_func_output(output, params=params)\n",
    "\n",
    "    mod = builder.get()\n",
    "\n",
    "    params = {\"linear_weight\": weight1_np, \"linear_weight1\": weight2_np}\n",
    "    mod = relax.transform.BindParams(\"main\", params)(mod)\n",
    "\n",
    "    target = \"llvm\"\n",
    "    return tvm.compile(mod, target)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd2b65c8",
   "metadata": {},
   "source": [
    "## 测试 conv2d cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f233a6d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name                          Duration (us)  Percent  Device  Count                                  Argument Shapes  \n",
      "matmul                                12.94    15.70    cpu0      2  float32[1, 64], float32[64, 64], float32[1, 64]  \n",
      "vm.builtin.check_tensor_info           1.61     1.95    cpu0      1                                   float32[1, 64]  \n",
      "relu                                   1.43     1.74    cpu0      2                   float32[1, 64], float32[1, 64]  \n",
      "vm.builtin.match_shape                 1.26     1.53    cpu0      1                                   float32[1, 64]  \n",
      "----------                                                                                                            \n",
      "Sum                                   17.23    20.92              6                                                   \n",
      "Total                                 82.39             cpu0      1                                                   \n",
      "\n",
      "Configuration\n",
      "-------------\n",
      "Number of threads: 24\n",
      "Executor: VM\n",
      "\n"
     ]
    }
   ],
   "source": [
    "data_np = np.random.randn(1, 64).astype(\"float32\")\n",
    "ex = get_exec(data_np.shape)\n",
    "\n",
    "vm = VirtualMachine(ex, tvm.cpu(), profile=True)\n",
    "report = vm.profile(\"main\", tvm.nd.array(data_np))\n",
    "print(report)\n",
    "\n",
    "assert \"Duration\" in str(report)\n",
    "assert \"matmul\" in str(report)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5337b2cf",
   "metadata": {},
   "source": [
    "## 测试 rpc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5674105",
   "metadata": {},
   "outputs": [],
   "source": [
    "def with_rpc(ex, f, data_np):\n",
    "    temp = utils.tempdir()\n",
    "    path = temp.relpath(\"vm_library.so\")\n",
    "    ex.export_library(path)\n",
    "\n",
    "    server = rpc.Server(\"127.0.0.1\")\n",
    "    remote = rpc.connect(server.host, server.port, session_timeout=10)\n",
    "\n",
    "    remote.upload(path)\n",
    "    rexec = remote.load_module(\"vm_library.so\")\n",
    "\n",
    "    device = remote.cpu()\n",
    "\n",
    "    vm = VirtualMachine(rexec, device=device, profile=True)\n",
    "    data = tvm.nd.array(data_np, device)\n",
    "\n",
    "    f(vm, data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e3d89c30",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-25 16:25:51.924 INFO bind to 127.0.0.1:9091\n",
      "2025-08-25 16:25:51.925 INFO connected from ('127.0.0.1', 36662)\n",
      "2025-08-25 16:25:51.926 INFO start serving at /tmp/tmptnsl0kx3\n",
      "2025-08-25 16:25:51.941 INFO load_module /tmp/tmptnsl0kx3/vm_library.so\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name                          Duration (us)  Percent  Device  Count                                  Argument Shapes  \n",
      "matmul                                 8.40    14.48    cpu0      2  float32[1, 64], float32[64, 64], float32[1, 64]  \n",
      "vm.builtin.check_tensor_info           1.72     2.97    cpu0      1                                   float32[1, 64]  \n",
      "relu                                   0.92     1.58    cpu0      2                   float32[1, 64], float32[1, 64]  \n",
      "vm.builtin.match_shape                 0.78     1.34    cpu0      1                                   float32[1, 64]  \n",
      "----------                                                                                                            \n",
      "Sum                                   11.81    20.37              6                                                   \n",
      "Total                                 58.00             cpu0      1                                                   \n",
      "\n",
      "Configuration\n",
      "-------------\n",
      "Number of threads: 24\n",
      "Executor: VM\n",
      "\n"
     ]
    }
   ],
   "source": [
    "data_np = np.random.randn(1, 64).astype(\"float32\")\n",
    "ex = get_exec(data_np.shape)\n",
    "\n",
    "def callback(vm, data):\n",
    "    vm.profile(\"main\", data)\n",
    "\n",
    "    vm.set_input(\"main\", data)\n",
    "    report = vm.profile(\"main\")\n",
    "\n",
    "    assert \"matmul\" in str(report)\n",
    "    print(report)\n",
    "\n",
    "with_rpc(ex, callback, data_np)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a107331",
   "metadata": {},
   "source": [
    "### 测试元组"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d178624b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-25 16:26:22.805 INFO bind to 127.0.0.1:9091\n",
      "2025-08-25 16:26:22.806 INFO connected from ('127.0.0.1', 33636)\n",
      "2025-08-25 16:26:22.806 INFO start serving at /tmp/tmpdvxl0sd8\n",
      "2025-08-25 16:26:22.821 INFO load_module /tmp/tmpdvxl0sd8/vm_library.so\n"
     ]
    }
   ],
   "source": [
    "@tvm.script.ir_module\n",
    "class NestedTuple:\n",
    "    @R.function\n",
    "    def main(\n",
    "        x: R.Tensor((16,), \"float32\")\n",
    "    ) -> R.Tuple(\n",
    "        R.Tuple(\n",
    "            R.Tensor((16,), \"float32\"),\n",
    "            R.Tuple(\n",
    "                R.Tensor((16,), \"float32\"),\n",
    "            ),\n",
    "        ),\n",
    "        R.Tensor((16,), \"float32\"),\n",
    "    ):\n",
    "        return ((x, (x,)), x)\n",
    "\n",
    "target = \"llvm\"\n",
    "ex = tvm.compile(NestedTuple, target)\n",
    "\n",
    "data_np = np.random.randn(16).astype(\"float32\")\n",
    "\n",
    "def callback(vm, data):\n",
    "    report = vm.profile(\"main\", data)\n",
    "    assert \"vm.builtin.make_tuple\" in str(report)\n",
    "\n",
    "with_rpc(ex, callback, data_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76f63992",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
